package net.siufung.core.utils;

import org.reactivestreams.Subscription;
import reactor.core.Disposable;
import reactor.core.publisher.BaseSubscriber;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;
import reactor.core.publisher.SignalType;
import reactor.core.scheduler.Scheduler;
import reactor.core.scheduler.Schedulers;
import reactor.util.function.Tuple2;

import javax.annotation.Nonnull;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.function.BiPredicate;

/**
 * @author 陈建峰
 * @since 2022/5/11 5:03 下午
 */
public class FluxUtil {

    /**
     * 缓冲比例
     * @param flux 流
     * @param rate 比例
     * @param maxTimeout 最大时间
     * @param <T> T
     * @return Flux<List<T>>
     */
    public static <T> Flux<List<T>> bufferRate(Flux<T> flux,
                                               int rate,
                                               Duration maxTimeout) {
        return bufferRate(flux, rate, 100, maxTimeout);
    }

    /**
     * 缓冲比例
     * @param flux 流
     * @param rate 比例
     * @param maxSize 最大缓冲大小
     * @param maxTimeout 最大时间
     * @param <T> T
     * @return Flux<List<T>>
     */
    public static <T> Flux<List<T>> bufferRate(Flux<T> flux,
                                               int rate,
                                               int maxSize,
                                               Duration maxTimeout) {
        return Flux.create(sink -> {
            BufferRateSubscriber<T> subscriber =
                    new BufferRateSubscriber<>(sink, maxSize, rate, maxTimeout, (e, arr) -> arr.size() >= maxSize);
            flux.elapsed().subscribe(subscriber);
            sink.onDispose(subscriber);
        });

    }

    /**
     * 缓冲比例
     * @param flux 流
     * @param rate 比例
     * @param maxSize 最大缓冲大小
     * @param maxTimeout 最大时间
     * @param flushCondition 条件
     * @param <T> T
     * @return Flux<List<T>>
     */
    public static <T> Flux<List<T>> bufferRate(Flux<T> flux,
                                               int rate,
                                               int maxSize,
                                               Duration maxTimeout,
                                               BiPredicate<T, List<T>> flushCondition) {
        return Flux.create(sink -> {
            BufferRateSubscriber<T> subscriber =
                    new BufferRateSubscriber<>(sink, maxSize, rate, maxTimeout, (e, arr) -> flushCondition
                    .test(e, arr) || arr.size() >= maxSize);
            flux.elapsed().subscribe(subscriber);
            sink.onDispose(subscriber);
        });
    }

    /**
     * 缓冲订阅比例
     * @param <T>
     */
    static class BufferRateSubscriber<T> extends BaseSubscriber<Tuple2<Long, T>> {
        /**
         * 缓冲大小
         */
        int bufferSize;
        /**
         * 比例
         */
        int rate;
        /**
         * 缓冲大小
         */
        volatile List<T> bufferArray;
        /**
         * 发射器
         */
        FluxSink<List<T>> fluxSink;
        /**
         * 时间
         */
        Duration timeout;
        /**
         * 调度器
         */
        Scheduler timer = Schedulers.parallel();
        /**
         * 中止Flux流
         */
        Disposable timerDispose;
        /**
         *
         */
        private final BiPredicate<T, List<T>> flushCondition;

        BufferRateSubscriber(FluxSink<List<T>> fluxSink, int bufferSize, int rate,
                             Duration timeout, BiPredicate<T, List<T>> flushCondition) {
            this.fluxSink = fluxSink;
            this.bufferSize = bufferSize;
            this.rate = rate;
            this.timeout = timeout;
            this.flushCondition = flushCondition;
            newBuffer();
        }

        protected List<T> newBuffer() {
            return bufferArray = new ArrayList<>(bufferSize);
        }

        @Override
        protected void hookFinally(@Nonnull SignalType signalType) {
            doFlush();
        }

        void doFlush() {
            if (bufferArray.size() > 0) {
                fluxSink.next(newBuffer());
            }
            request(bufferSize);
            if (timerDispose != null && !timerDispose.isDisposed()) {
                timerDispose.dispose();
            }
        }

        @Override
        protected void hookOnSubscribe(@Nonnull Subscription subscription) {
            request(bufferSize);
        }

        @Override
        protected void hookOnNext(Tuple2<Long, T> value) {
            bufferArray.add(value.getT2());
            if (value.getT1() > rate) {
                doFlush();
            } else {
                if (flushCondition.test(value.getT2(), bufferArray)) {
                    doFlush();
                } else {
                    if (timerDispose == null || timerDispose.isDisposed()) {
                        timerDispose = timer.schedule(this::doFlush, timeout.toMillis(), TimeUnit.MILLISECONDS);
                    }
                }
            }
        }
    }
}
